import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.spatial.transform import Rotation

def plot_tum_trajectory_with_pose_axes(filename):
    """
    从TUM文件中读取三维点及旋转信息并绘制姿态坐标轴。

    参数:
    filename -- TUM格式的文件名
    """
    # 读取文件
    data = np.loadtxt(filename, delimiter=' ', usecols=(1, 2, 3, 4, 5, 6, 7))

    # 提取x, y, z坐标和四元数
    x = data[:, 0]
    y = data[:, 1]
    z = data[:, 2]
    qx = data[:, 3]
    qy = data[:, 4]
    qz = data[:, 5]
    qw = data[:, 6]
    
    # 创建3D绘图
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 绘制轨迹
    ax.plot(x, y, z, label='Trajectory', color='b')
    ax.scatter(x, y, z, c='b', marker='o', s=0.1)

    # # 绘制每个位姿的坐标轴
    # for i in range(len(x)):
    #     # 使用四元数创建旋转对象
    #     rotation = Rotation.from_quat([qx[i], qy[i], qz[i], qw[i]])

    #     # 定义坐标轴方向
    #     axes = np.array([[1, 0, 0],   # X轴
    #                      [0, 1, 0],   # Y轴
    #                      [0, 0, 1]])  # Z轴

    #     # 旋转坐标轴方向
    #     rotated_axes = rotation.apply(axes)

    #     # 绘制坐标轴
    #     ax.quiver(x[i], y[i], z[i], rotated_axes[0, 0], rotated_axes[0, 1], rotated_axes[0, 2], 
    #                color='r', length=0.1)  # X轴
    #     ax.quiver(x[i], y[i], z[i], rotated_axes[1, 0], rotated_axes[1, 1], rotated_axes[1, 2], 
    #                color='g', length=0.1)  # Y轴
    #     ax.quiver(x[i], y[i], z[i], rotated_axes[2, 0], rotated_axes[2, 1], rotated_axes[2, 2], 
    #                color='b', length=0.1)  # Z轴

    # 手动设置坐标轴的尺度相同
    max_range = np.array([x.max()-x.min(), y.max()-y.min(), z.max()-z.min()]).max() / 2.0
    mid_x = (x.max()+x.min()) * 0.5
    mid_y = (y.max()+y.min()) * 0.5
    mid_z = (z.max()+z.min()) * 0.5

    ax.set_xlim(mid_x - max_range, mid_x + max_range)
    ax.set_ylim(mid_y - max_range, mid_y + max_range)
    ax.set_zlim(mid_z - max_range, mid_z + max_range)

    # 设置标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title('TUM Trajectory with Pose Axes')

    ax.legend()

    # 显示图形
    plt.show()


data_dir = '../log/A1/park/'
file_name = 'LIOSAM_XYZ_align.txt'

data_file = data_dir + file_name
print("3D Path Show is Start")
plot_tum_trajectory_with_pose_axes(data_file)